# py/sdk.py
import qianfan
import os
import re

class QianfanWrapper:
    """优化的千帆调用封装（支持超时和重试）"""
    
    def __init__(self):
        self._init_qianfan()
        self.model = "ERNIE-3.5-8K"
        self.max_retries = 3  # 最大重试次数
    
    def _init_qianfan(self):
        """初始化认证（从环境变量读取更安全）"""
        os.environ["QIANFAN_ACCESS_KEY"] = os.getenv("QIANFAN_AK", "ALTAKiwnHEggAvYXbrd21K3luc")
        os.environ["QIANFAN_SECRET_KEY"] = os.getenv("QIANFAN_SK", "0d9cd08f6a924dac8673281d172ba485")
    
    def query(self, prompt: str, timeout: int = 5) -> str:
        """调用千帆API（自动重试+超时控制）"""
        chat = qianfan.ChatCompletion(model=self.model)
        
        for attempt in range(self.max_retries):
            try:
                resp = chat.do(
                    messages=[{"role": "user", "content": prompt}],
                    stream=False,
                    request_timeout=timeout  # 千坊SDK实际参数名
                )
                return self._clean_response(resp['result'])
            except Exception as e:
                if attempt == self.max_retries - 1:
                    raise
                print(f"第{attempt+1}次尝试失败: {str(e)}")

    def _clean_response(self, text: str) -> list:
         # 1. 去除首尾多余的方括号和引号
        cleaned = text.strip().strip('[]"\'')
        
        # 2. 分割逗号分隔的带引号关键词
        keywords = [kw.strip(' "\'') for kw in cleaned.split(',') if kw.strip()]
        
        return keywords[:20]  # 限制数量

def extract_entities(question: str) -> str:
    qianfan = QianfanWrapper()
    
    # 多轮解析prompt
    prompt = f"""
    请分析以下查询并提取：
    1. 查询主体（人物/地点/组织）
    2. 查询方向（关系/属性/事件）
    3. 所有关键实体
    4.输出要求：直接返回纯文本关键词列表，用空格分隔
    示例：
    输入："刘备三兄弟的仕途"
    输出：["刘备","关羽","张飞","仕途"]

    输入："赤壁之战参战方"
    输出：["赤壁之战","曹操","孙权","刘备"]
 
    现在请处理：
    输入：{question}
    输出："""
    try:
        result = qianfan.query(prompt, timeout=10)
        # 确保返回的是列表（即使为空）
        return result if isinstance(result, list) else []
    except Exception as e:
        print(f"Error: {str(e)}")
        return []  # 错误时返回空数组